package com.baidu.acg.industry.bigdata.controller;

import com.baidu.acg.industry.bigdata.common.CommonResponse;
import com.baidu.acg.industry.bigdata.controller.request.ReasonRequest;
import com.baidu.acg.industry.bigdata.controller.request.TrainModelRequest;
import com.baidu.acg.industry.bigdata.controller.request.UpdateModelParamsRequest;
import com.baidu.acg.industry.bigdata.controller.response.FilePathResponse;
import com.baidu.acg.industry.bigdata.controller.response.PreEpochResponse;
import com.baidu.acg.industry.bigdata.controller.response.PreSubmitResponse;
import com.baidu.acg.industry.bigdata.controller.response.ResponseData;
import com.baidu.acg.industry.bigdata.service.SiFangModelService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.PutMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

@RestController
@RequestMapping("/v1/platform/sifang")
@Slf4j
public class SiFangPlatFormController {
    private static final AtomicInteger PRE_COUNTER = new AtomicInteger(0);
    private static final AtomicInteger TRAIN_COUNTER = new AtomicInteger(0);

    @Autowired
    private SiFangModelService siFangModelService;

    // TODO 需要读取txt文件取模型的物理参数

    /**
     * 提交预处理任务
     */
    @PostMapping("/data/pre")
    public ResponseEntity<PreSubmitResponse> runPreTask(@RequestParam("modelTaskId") String modelTaskId,
                                                        @RequestParam("dataSetPath") String dataSetPath,
                                                        @RequestParam("dataSetName") String dataSetName) {
        // 解析txt
//        siFangModelService.preJob(dataSetPath, modelTaskId);
        // 首先根据传递过来的txt文件，然后解析出来
        PreSubmitResponse preEpochResponse = new PreSubmitResponse();
        preEpochResponse.setModelTaskId("pre_1_eax13131");
        return CommonResponse.ok(preEpochResponse);
    }

    /**
     * 修改模型yaml文件
     */
    @GetMapping("/modify/yaml")
    public Map<String, List<String>> modifyYaml() {
//        siFangModelService.modifyModelYaml();
        Map<String, List<String>> map = new HashMap<>();
        List<String> filePath = map.put("filePath", Arrays.asList("123", "123"));

        return map;
    }

    /**
     * 预处理进度
     * modelTaskId:pre_taskId_uuid
     */

    @GetMapping("/pre/epoch")
    public ResponseEntity<PreEpochResponse> preDealEpoch(@RequestParam("modelTaskId") String modelTaskId) {

        int count = PRE_COUNTER.updateAndGet(n -> (n >= 3) ? 0 : n + 1); // 循环计数
        PreEpochResponse preEpochResponse = new PreEpochResponse();
        switch (count) {
            case 1:
                preEpochResponse = siFangModelService.getPreDataStream(modelTaskId, "json", "pre");
                return CommonResponse.ok(preEpochResponse);
            case 2:
                preEpochResponse.setEpoch("30");
                return CommonResponse.ok(preEpochResponse);
            case 3:
                preEpochResponse.setEpoch("60");
                return CommonResponse.ok(preEpochResponse);
            default:
                preEpochResponse.setEpoch("100");
                PRE_COUNTER.set(0); // 重置计数器
                return CommonResponse.ok(preEpochResponse);
        }
    }


    /**
     * 预处理日志
     */

    @GetMapping("/pre/log")
    public ResponseEntity<FilePathResponse> preDealFileLog(@RequestParam("modelTaskId") String modelTaskId) {
        FilePathResponse response = siFangModelService.getModelFileList(modelTaskId, "log");
        return CommonResponse.ok(response);
    }

    /**
     * 预处理数据
     *
     * @param modelTaskId
     * @return
     */
    @GetMapping("/pre/dataset")
    public ResponseEntity<FilePathResponse> preDealFile(@RequestParam("modelTaskId") String modelTaskId) {
        FilePathResponse response = siFangModelService.getModelFileList(modelTaskId, "data");
        return CommonResponse.ok(response);

    }

    /**
     * @param taskId
     * @return
     */
    @PostMapping("/server/stop")
    public ResponseEntity<ResponseData> stopPreTask(@RequestParam("taskId") String taskId) {
        ResponseData responseData = new ResponseData();
        responseData.setMessage("success");
        return CommonResponse.ok(responseData);
    }


    @PostMapping("/train")
    public ResponseEntity<PreSubmitResponse> runTrainTask(@RequestBody TrainModelRequest trainModelRequest) {
        // 解析txt
        log.info("提交训练任务=>" + trainModelRequest);
//        siFangModelService.preJob(dataSetPath, modelTaskId);
        // 首先根据传递过来的txt文件，然后解析出来
        PreSubmitResponse preEpochResponse = new PreSubmitResponse();
        preEpochResponse.setModelTaskId("train_1_eax13131");
        return CommonResponse.ok(preEpochResponse);
    }

    /**
     * 项目获取训练日志
     */
    @GetMapping("/train/log")
    public ResponseEntity<FilePathResponse> trainLogFile(@RequestParam("modelTaskId") String modelTaskId) {
        FilePathResponse response = siFangModelService.getModelFileList(modelTaskId, "log");
        return CommonResponse.ok(response);
    }


    /**
     * 获取训练进度
     */
    @GetMapping("/train/process")
    public ResponseEntity<PreEpochResponse> trainProcess(@RequestParam("modelTaskId") String modelTaskId) {

        int count = TRAIN_COUNTER.updateAndGet(n -> (n >= 3) ? 0 : n + 1); // 循环计数
        PreEpochResponse preEpochResponse = new PreEpochResponse();
        switch (count) {
            case 1:
                preEpochResponse = siFangModelService.getPreDataStream(modelTaskId, "json", "train");
                return CommonResponse.ok(preEpochResponse);
            case 2:
                preEpochResponse.setEpoch("30");
                return CommonResponse.ok(preEpochResponse);
            case 3:
                preEpochResponse.setEpoch("40");
                return CommonResponse.ok(preEpochResponse);
            default:
                preEpochResponse.setEpoch("100");
                TRAIN_COUNTER.set(0); // 重置计数器
                return CommonResponse.ok(preEpochResponse);
        }
    }

    /**
     * 训练数据
     *
     * @param modelTaskId
     * @return
     */
    @GetMapping("/train/list")
    public ResponseEntity<FilePathResponse> trainDataFile(@RequestParam("modelTaskId") String modelTaskId) {
        FilePathResponse filePathResponse = siFangModelService.getModelFileList(modelTaskId, "json");
        return CommonResponse.ok(filePathResponse);
    }

    /**
     * 训练vtp数据
     *
     * @param modelTaskId
     * @return
     */
    @GetMapping("/train/vtp")
    public ResponseEntity<FilePathResponse> trainDataVtpFile(@RequestParam("modelTaskId") String modelTaskId) {
        FilePathResponse filePathResponse = siFangModelService.getModelFileList(modelTaskId, "vtp");
        return CommonResponse.ok(filePathResponse);
    }

    @PostMapping("/train/stop")
    public ResponseEntity<ResponseData> stopTrainTask(@RequestParam("taskId") String taskId) {
        ResponseData responseData = new ResponseData();
        responseData.setMessage("success");
        return CommonResponse.ok(responseData);
    }


    /**
     * 查询推理接⼝数据
     *
     */
    @GetMapping("/reason/model")
    public ResponseEntity<FilePathResponse> queryReasonModel(@RequestParam("modelTaskId") String modelTaskId, @RequestParam("objectName") String objectName) {
        FilePathResponse filePathResponse = siFangModelService.getModelFileList(modelTaskId, "json");
        return CommonResponse.ok(filePathResponse);
    }

    /**
     * 查询推理vtp文件接⼝
     *
     */
    @GetMapping("/reason/vtp")
    public ResponseEntity<FilePathResponse> queryReasonModelVtp(@RequestParam("modelTaskId") String modelTaskId) {
        FilePathResponse filePathResponse = siFangModelService.getModelFileList(modelTaskId, "vtp");
        return CommonResponse.ok(filePathResponse);
    }

    /**
     * 查询推理⽇志
     *
     */
    @GetMapping("/reason/log")
    public ResponseEntity<FilePathResponse> reasonLogFile(@RequestParam("modelTaskId") String modelTaskId) {
        FilePathResponse response = siFangModelService.getModelFileList(modelTaskId, "log");
        return CommonResponse.ok(response);
    }


    /**
     * 服务上线
     */
    @PostMapping("/model/online")
    public ResponseEntity<ResponseData> onlineModelServer(@RequestParam("modelTaskId") String modelTaskId,
                                                          @RequestParam("gpuNum") String gpuNum) {
        ResponseData responseData = new ResponseData();
        responseData.setMessage("success");
        return CommonResponse.ok(responseData);
    }


    /**
     * 服务下线
     */
    @PostMapping("/model/offline")
    public ResponseEntity<ResponseData> offlineModelServer(@RequestParam("modelTaskId") String modelTaskId,
                                                           @RequestParam("gpuNum") String gpuNum) {
        ResponseData responseData = new ResponseData();
        responseData.setMessage("success");
        return CommonResponse.ok(responseData);
    }


    @GetMapping("/reason/run")
    public ResponseEntity<ResponseData> startReason(@RequestBody ReasonRequest reasonRequest) {
        ResponseData responseData = new ResponseData();
        responseData.setEndpoint("success");
        return CommonResponse.ok(responseData);
    }

    /**
     * 推理服务查询接口
     */
    @GetMapping("/reason/server")
    public ResponseEntity<ResponseData> fflineModelServer(@RequestParam("modelTaskId") String modelTaskId,
                                                          @RequestParam("gpuNum") String gpuNum) {

        ResponseData responseData = new ResponseData();
        responseData.setEndpoint("http://127.0.0.1:8089" + "/v1/platform/sifang/reason/run");
        return CommonResponse.ok(responseData);
    }


    /**
     * 修改模型训练参数接口
     */
    @PutMapping("/update/model/params")
    public ResponseEntity<ResponseData> updateModelParams(@RequestBody UpdateModelParamsRequest updateModelParamsRequest) {
        ResponseData responseData = new ResponseData();
        responseData.setEndpoint("http://127.0.0.1:8089" + "/v1/platform/sifang/reason/run");
        return CommonResponse.ok(responseData);
    }


    /**
     * 任务状态查询接口
     */
    @GetMapping("/server/task/{workflowName}")
    public ResponseEntity<ResponseData> getServerStatus(@PathVariable("workflowName") String workflowName) {
        ResponseData responseData = new ResponseData();
        responseData.setMessage("success");
        return CommonResponse.ok(responseData);
    }


}


